#!/bin/bash

# 5GC Testbed Comprehensive Test Runner
# This script runs all test scenarios for the 5GC testbed

set -e  # Exit on any error

# Configuration
VM1="192.168.1.206"  # Core Network
VM2="192.168.1.219"  # UPF
VM3="192.168.1.221"  # gNodeB
VM4="192.168.1.205"  # UE

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Test results
declare -A test_results
declare -A test_messages

# Logging function
log() {
    echo -e "${GREEN}[$(date +'%Y-%m-%d %H:%M:%S')] $1${NC}"
}

warn() {
    echo -e "${YELLOW}[$(date +'%Y-%m-%d %H:%M:%S')] WARNING: $1${NC}"
}

error() {
    echo -e "${RED}[$(date +'%Y-%m-%d %H:%M:%S')] ERROR: $1${NC}"
}

info() {
    echo -e "${BLUE}[$(date +'%Y-%m-%d %H:%M:%S')] INFO: $1${NC}"
}

# Function to record test result
record_test() {
    local test_name=$1
    local result=$2
    local message=$3
    
    test_results[$test_name]=$result
    test_messages[$test_name]=$message
}

# Function to check if a service is running
check_service() {
    local service=$1
    local vm=$2
    
    if ssh -o ConnectTimeout=5 $vm "systemctl is-active --quiet $service" 2>/dev/null; then
        return 0
    else
        return 1
    fi
}

# Function to check if a process is running
check_process() {
    local process=$1
    local vm=$2
    
    if ssh -o ConnectTimeout=5 $vm "pgrep -f '$process'" >/dev/null 2>&1; then
        return 0
    else
        return 1
    fi
}

# Function to check network connectivity
check_connectivity() {
    local vm=$1
    
    if ping -c 1 -W 5 $vm >/dev/null 2>&1; then
        return 0
    else
        return 1
    fi
}

# Function to check port connectivity
check_port() {
    local vm=$1
    local port=$2
    local protocol=${3:-tcp}
    
    if [ "$protocol" = "udp" ]; then
        if nc -zu $vm $port 2>/dev/null; then
            return 0
        else
            return 1
        fi
    else
        if nc -z $vm $port 2>/dev/null; then
            return 0
        else
            return 1
        fi
    fi
}

# Print test header
print_header() {
    echo
    echo "=========================================="
    echo "        5GC Testbed Test Suite"
    echo "=========================================="
    echo "Date: $(date)"
    echo "=========================================="
    echo
}

# Print test results
print_results() {
    echo
    echo "=========================================="
    echo "           Test Results Summary"
    echo "=========================================="
    echo
    
    local total_tests=0
    local passed_tests=0
    local failed_tests=0
    local warning_tests=0
    
    for test_name in "${!test_results[@]}"; do
        total_tests=$((total_tests + 1))
        
        case ${test_results[$test_name]} in
            "PASS")
                echo -e "✓ ${GREEN}$test_name${NC}: ${test_messages[$test_name]}"
                passed_tests=$((passed_tests + 1))
                ;;
            "FAIL")
                echo -e "✗ ${RED}$test_name${NC}: ${test_messages[$test_name]}"
                failed_tests=$((failed_tests + 1))
                ;;
            "WARN")
                echo -e "⚠ ${YELLOW}$test_name${NC}: ${test_messages[$test_name]}"
                warning_tests=$((warning_tests + 1))
                ;;
        esac
    done
    
    echo
    echo "=========================================="
    echo "Test Summary:"
    echo "  Total Tests: $total_tests"
    echo -e "  ${GREEN}Passed: $passed_tests${NC}"
    echo -e "  ${YELLOW}Warnings: $warning_tests${NC}"
    echo -e "  ${RED}Failed: $failed_tests${NC}"
    echo "=========================================="
    
    if [ $failed_tests -gt 0 ]; then
        exit 1
    fi
}

# Test 1: Network Connectivity
test_network_connectivity() {
    log "Running Network Connectivity Tests..."
    
    for vm in $VM1 $VM2 $VM3 $VM4; do
        if check_connectivity $vm; then
            record_test "Network_$vm" "PASS" "VM $vm is reachable"
        else
            record_test "Network_$vm" "FAIL" "VM $vm is not reachable"
        fi
    done
}

# Test 2: Service Health
test_service_health() {
    log "Running Service Health Tests..."
    
    # Core services on VM1
    services=("open5gs-amfd" "open5gs-smfd" "open5gs-udmd" "open5gs-udrd" "open5gs-pcfd")
    
    for service in "${services[@]}"; do
        if check_service $service $VM1; then
            record_test "$service" "PASS" "Service is running"
        else
            record_test "$service" "FAIL" "Service is not running"
        fi
    done
    
    # UPF service on VM2
    if check_service "open5gs-upfd" $VM2; then
        record_test "open5gs-upfd" "PASS" "UPF service is running"
    else
        record_test "open5gs-upfd" "FAIL" "UPF service is not running"
    fi
}

# Test 3: UERANSIM Processes
test_ueransim_processes() {
    log "Running UERANSIM Process Tests..."
    
    if check_process "nr-gnb" $VM3; then
        record_test "gNodeB_Process" "PASS" "gNodeB process is running"
    else
        record_test "gNodeB_Process" "FAIL" "gNodeB process is not running"
    fi
    
    if check_process "nr-ue" $VM4; then
        record_test "UE_Process" "PASS" "UE process is running"
    else
        record_test "UE_Process" "FAIL" "UE process is not running"
    fi
}

# Test 4: Port Connectivity
test_port_connectivity() {
    log "Running Port Connectivity Tests..."
    
    # Core network ports (VM1)
    if check_port $VM1 7777; then
        record_test "Core_SBI_Port" "PASS" "SBI port 7777 is listening"
    else
        record_test "Core_SBI_Port" "FAIL" "SBI port 7777 is not listening"
    fi
    
    if check_port $VM1 38412; then
        record_test "Core_NGAP_Port" "PASS" "NGAP port 38412 is listening"
    else
        record_test "Core_NGAP_Port" "FAIL" "NGAP port 38412 is not listening"
    fi
    
    # UPF ports (VM2)
    if check_port $VM2 8805; then
        record_test "UPF_PFCP_Port" "PASS" "UPF PFCP port 8805 is listening"
    else
        record_test "UPF_PFCP_Port" "FAIL" "UPF PFCP port 8805 is not listening"
    fi
    
    if check_port $VM2 2152 udp; then
        record_test "UPF_GTPU_Port" "PASS" "UPF GTP-U port 2152 is listening"
    else
        record_test "UPF_GTPU_Port" "FAIL" "UPF GTP-U port 2152 is not listening"
    fi
}

# Test 5: SCTP Associations
test_sctp_associations() {
    log "Running SCTP Association Tests..."
    
    # Check gNodeB-AMF association
    if ssh $VM3 "ss -tan | grep '192.168.1.10:38412' | grep ESTAB" >/dev/null 2>&1; then
        record_test "gNB_AMF_SCTP" "PASS" "gNodeB-AMF SCTP association established"
    else
        record_test "gNB_AMF_SCTP" "FAIL" "gNodeB-AMF SCTP association not established"
    fi
    
    # Check SMF-UPF association
    if ssh $VM1 "ss -tan | grep '192.168.1.11:8805' | grep ESTAB" >/dev/null 2>&1; then
        record_test "SMF_UPF_SCTP" "PASS" "SMF-UPF SCTP association established"
    else
        record_test "SMF_UPF_SCTP" "FAIL" "SMF-UPF SCTP association not established"
    fi
}

# Test 6: Database Connectivity
test_database_connectivity() {
    log "Running Database Connectivity Tests..."
    
    if ssh $VM1 "mongo --eval 'db.runCommand(\"ping\")' --quiet" >/dev/null 2>&1; then
        record_test "MongoDB_Connectivity" "PASS" "MongoDB is accessible"
    else
        record_test "MongoDB_Connectivity" "FAIL" "MongoDB is not accessible"
    fi
}

# Test 7: UE Registration
test_ue_registration() {
    log "Running UE Registration Tests..."
    
    # Check if UE is registered by looking for registration success in logs
    if ssh $VM4 "grep -i 'registration successful' ue.log | tail -1" >/dev/null 2>&1; then
        record_test "UE_Registration" "PASS" "UE is registered"
    else
        record_test "UE_Registration" "FAIL" "UE is not registered"
    fi
}

# Test 8: PDU Session
test_pdu_session() {
    log "Running PDU Session Tests..."
    
    # Check if PDU session is established
    if ssh $VM4 "grep -i 'pdu session established' ue.log | tail -1" >/dev/null 2>&1; then
        record_test "PDU_Session" "PASS" "PDU session is established"
    else
        record_test "PDU_Session" "FAIL" "PDU session is not established"
    fi
}

# Test 9: Performance Tests
test_performance() {
    log "Running Performance Tests..."
    
    # Test latency
    for vm in $VM1 $VM2 $VM3 $VM4; do
        avg_latency=$(ping -c 5 $vm | tail -1 | awk '{print $4}' | cut -d'/' -f2 2>/dev/null || echo "0")
        if (( $(echo "$avg_latency < 10" | bc -l 2>/dev/null || echo "1") )); then
            record_test "Latency_$vm" "PASS" "Latency: ${avg_latency}ms"
        elif (( $(echo "$avg_latency < 50" | bc -l 2>/dev/null || echo "1") )); then
            record_test "Latency_$vm" "WARN" "High latency: ${avg_latency}ms"
        else
            record_test "Latency_$vm" "FAIL" "Very high latency: ${avg_latency}ms"
        fi
    done
}

# Test 10: System Resources
test_system_resources() {
    log "Running System Resource Tests..."
    
    # Check CPU usage
    for vm in $VM1 $VM2 $VM3 $VM4; do
        cpu_usage=$(ssh $vm "top -bn1 | grep 'Cpu(s)' | awk '{print \$2}' | cut -d'%' -f1" 2>/dev/null || echo "0")
        if (( $(echo "$cpu_usage < 80" | bc -l 2>/dev/null || echo "1") )); then
            record_test "CPU_$vm" "PASS" "CPU usage: ${cpu_usage}%"
        elif (( $(echo "$cpu_usage < 90" | bc -l 2>/dev/null || echo "1") )); then
            record_test "CPU_$vm" "WARN" "High CPU usage: ${cpu_usage}%"
        else
            record_test "CPU_$vm" "FAIL" "Very high CPU usage: ${cpu_usage}%"
        fi
    done
    
    # Check memory usage
    for vm in $VM1 $VM2 $VM3 $VM4; do
        mem_usage=$(ssh $vm "free | grep Mem | awk '{printf \"%.1f\", \$3/\$2 * 100.0}'" 2>/dev/null || echo "0")
        if (( $(echo "$mem_usage < 80" | bc -l 2>/dev/null || echo "1") )); then
            record_test "Memory_$vm" "PASS" "Memory usage: ${mem_usage}%"
        elif (( $(echo "$mem_usage < 90" | bc -l 2>/dev/null || echo "1") )); then
            record_test "Memory_$vm" "WARN" "High memory usage: ${mem_usage}%"
        else
            record_test "Memory_$vm" "FAIL" "Very high memory usage: ${mem_usage}%"
        fi
    done
}

# Main test function
main() {
    print_header
    
    # Run all tests
    test_network_connectivity
    test_service_health
    test_ueransim_processes
    test_port_connectivity
    test_sctp_associations
    test_database_connectivity
    test_ue_registration
    test_pdu_session
    test_performance
    test_system_resources
    
    # Print results
    print_results
}

# Handle command line arguments
case "${1:-}" in
    "--help"|"-h")
        echo "Usage: $0 [OPTIONS]"
        echo
        echo "Options:"
        echo "  -h, --help          Show this help message"
        echo
        echo "Examples:"
        echo "  $0                  Run all tests"
        echo
        echo "Test Categories:"
        echo "  - Network Connectivity"
        echo "  - Service Health"
        echo "  - UERANSIM Processes"
        echo "  - Port Connectivity"
        echo "  - SCTP Associations"
        echo "  - Database Connectivity"
        echo "  - UE Registration"
        echo "  - PDU Session"
        echo "  - Performance"
        echo "  - System Resources"
        ;;
    *)
        main
        ;;
esac 